import numpy.matlib
import numpy as np
import pandas as pd

from .factor_model import FactorModel
from .returns import Returns


class Portfolio(object):
    """
    Gather asset weights or factor exposures automatically.

    The indexes of weights and exposures should be fund_id.
    """

    def __init__(self) -> None:
        self.asset_returns = {}  # save returns
        self.weights = pd.DataFrame()
        self.alpha = None
        self.indivial_risk = None
        self.exposures = None
        return None

    def add_asset(self, asset_return: Returns, weight: float = 0) -> None:
        """Add asset to portfolio and set weight (optional).

        Args:
            asset_return (Returns): asset return
            weight (float, optional): weight. Defaults to 0.
        """
        name = asset_return.name
        self.asset_returns[name] = asset_return
        tmp_df = pd.DataFrame({"weights": {name: weight}})
        self.weights = pd.concat([self.weights, tmp_df])
        return None

    def analyse_assets(self, factor_model: FactorModel, alpha: bool = True) -> None:
        """Analyse asset returns using the inputted factor model.

        Args:
            factor_model (FactorModel): factor model to be used
        """
        self.alpha = pd.DataFrame()
        self.indivial_risk = pd.DataFrame()
        self.exposures = pd.DataFrame()
        for asset in self.asset_returns.keys():
            result = factor_model.get_asset_exposure(
                self.asset_returns[asset], alpha=alpha
            )
            self.alpha = pd.concat(
                [self.alpha, pd.DataFrame({"alpha": {asset: result["alpha"]}})]
            )
            self.indivial_risk = pd.concat(
                [
                    self.indivial_risk,
                    pd.DataFrame({"indivial_risk": {asset: result["indivial_risk"]}}),
                ]
            )
            self.exposures = pd.concat(
                [self.exposures, pd.DataFrame({asset: result["exposure"]}).T]
            )
            self.exposures = self.exposures[factor_model.factor_list]
        return None

    def get_portfolio_exposures(self) -> dict:
        all_data = pd.concat(
            [self.weights, self.alpha, self.indivial_risk, self.exposures],
            join="inner",
            axis=1,
        )
        w = all_data["weights"].values.reshape([-1, 1])
        alpha = all_data["alpha"].values.reshape([-1, 1])
        Epsilon = np.diagflat(all_data["indivial_risk"] ** 2)
        X = all_data[self.exposures.columns].values
        results = {}
        results["alpha"] = (w.T @ alpha)[0, 0]
        results["indivial_risk"] = np.sqrt((w.T @ Epsilon @ w)[0, 0])
        X_p = w.T @ X
        results["exposure"] = dict(zip(self.exposures.columns, X_p[0]))
        return results
